import torch
class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv = torch.nn.Sequential(
            # 1.卷积操作-卷积层
            torch.nn.Conv2d(1, 32, 5, 2),
            # 2.归一化BN层
            torch.nn.BatchNorm2d(32),
            # 3.Relu函数- 激活层
            torch.nn.ReLU(),
            # 4.最大池化-池化层
            torch.nn.MaxPool2d(kernel_size=2)
        )
        # 5.全连接层
        self.fc = torch.nn.Linear( 14*14*32, 10)
    def forward(self, out):
        out = self.conv(out)
        # 将图像展开成一维
        out = out.view(out.size()[0], -1)
        out = self.fc(out)
        return out